#!/usr/bin/env python3
"""
数据格式转换智能体
负责将数据分析结果转换为标准的前端可视化格式
"""

import json
import re
from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from config import DS_CONFIG, QWEN_CONFIG
from formatter_prompts import formatter_templates

# 导入LLM配置
try:
    from langchain_deepseek import ChatDeepSeek
    from langchain_openai import ChatOpenAI
    from langchain.schema import HumanMessage
except ImportError:
    print("⚠️ 缺少langchain依赖，请安装: pip install langchain langchain-community")

@dataclass
class FormatterInput:
    """格式转换智能体的输入数据"""
    user_request: str
    query_results: List[Dict[str, Any]]
    summary: str

@dataclass
class FormatterOutput:
    """格式转换智能体的输出数据"""
    summary: str
    chart_type: str
    data: Optional[Dict[str, Any]]

class DataFormatterAgent:
    """数据格式转换智能体"""
    
    def __init__(self, config=DS_CONFIG):
        """初始化格式转换智能体"""
        self.config = config
        self.llm = self._init_llm()
        
    def _init_llm(self):
        """初始化LLM模型"""
        try:
            if self.config.MODEL == 'deepseek-v3':
                return ChatDeepSeek(
                    model="deepseek-v3",
                    base_url=self.config.BASE_URL,
                    api_key=self.config.ALIYUN_API_KEY,
                    api_base=self.config.BASE_URL,
                    temperature=0.3,  # 降低温度以获得更稳定的输出
                    streaming=False
                )
            elif self.config.MODEL == 'qwen-max':
                return ChatOpenAI(
                    model="qwen-max",
                    base_url=self.config.BASE_URL,
                    api_key=self.config.ALIYUN_API_KEY,
                    temperature=0.3,
                    streaming=False
                )
            else:
                print(f"⚠️ 不支持的模型: {self.config.MODEL}")
                return None
        except Exception as e:
            print(f"❌ LLM初始化失败: {e}")
            return None
    
    def format_data(self, formatter_input: FormatterInput) -> FormatterOutput:
        """
        将数据分析结果转换为标准格式
        
        Args:
            formatter_input: 包含用户请求、查询结果和分析总结的输入数据
            
        Returns:
            FormatterOutput: 包含summary、chart_type和data的标准格式输出
        """
        try:
            print("🔄 开始数据格式转换...")
            
            # 检查输入数据
            if not formatter_input.query_results:
                print("⚠️ 查询结果为空，返回默认格式")
                return FormatterOutput(
                    summary=formatter_input.summary,
                    chart_type="",
                    data=None
                )
            
            # 构建LLM输入
            llm_input = self._build_llm_input(formatter_input)
        
            
            # 调用LLM进行格式转换
            if self.llm:
                formatted_result = self._call_llm_format(llm_input)
                if formatted_result:
                    return formatted_result
            
            # 如果LLM调用失败，使用备用逻辑
            print("⚠️ LLM调用失败，使用备用格式转换逻辑")
            return self._fallback_format(formatter_input)
            
        except Exception as e:
            print(f"❌ 数据格式转换失败: {e}")
            return FormatterOutput(
                summary=formatter_input.summary,
                chart_type="bar",
                data=None
            )
    
    def _build_llm_input(self, formatter_input: FormatterInput) -> str:
        """构建LLM输入"""
        # 准备数据预览（前5条记录）
        data_preview = formatter_input.query_results[:5]
        
        # 分析数据结构
        if data_preview:
            first_row = data_preview[0]
            fields_info = []
            for key, value in first_row.items():
                field_type = "数值" if isinstance(value, (int, float)) else "文本"
                fields_info.append(f"{key}({field_type})")
        else:
            fields_info = ["无数据"]
        
        # 构建输入数据
        input_data = {
            "user_request": formatter_input.user_request,
            "query_results": data_preview,
            "summary": formatter_input.summary
        }
        
        # 构建完整的prompt
        prompt = f"""
{formatter_templates}

## 当前输入数据
用户请求: {formatter_input.user_request}
数据字段: {', '.join(fields_info)}
数据行数: {len(formatter_input.query_results)}
数据预览: {json.dumps(data_preview, ensure_ascii=False, indent=2)}
现有总结: {formatter_input.summary}

请根据上述数据，生成标准格式的输出。
"""
        return prompt
    
    def _call_llm_format(self, prompt: str) -> Optional[FormatterOutput]:
        """调用LLM进行格式转换"""
        try:
            print("🤖 调用LLM进行格式转换...")
            
            # 调用LLM
            response = self.llm.invoke([HumanMessage(content=prompt)])
            content = response.content
            
            # 提取JSON
            json_result = self._extract_json(content)
            if json_result:
                # 验证输出格式
                if self._validate_output(json_result):
                    print("✅ LLM格式转换成功")
                    return FormatterOutput(
                        summary=json_result.get("summary", ""),
                        chart_type=json_result.get("chart_type", "bar"),
                        data=json_result.get("data")
                    )
            
            print("⚠️ LLM返回格式不正确")
            return None
            
        except Exception as e:
            print(f"❌ LLM调用失败: {e}")
            return None
    
    def _extract_json(self, content: str) -> Optional[Dict[str, Any]]:
        """从LLM响应中提取JSON"""
        try:
            # 移除markdown代码块标记
            content = re.sub(r'```json\s*', '', content)
            content = re.sub(r'```\s*', '', content)
            
            # 查找JSON对象
            match = re.search(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', content, re.DOTALL)
            if match:
                json_str = match.group(0)
                return json.loads(json_str)
            
            return None
            
        except Exception as e:
            print(f"❌ JSON提取失败: {e}")
            return None
    
    def _validate_output(self, output: Dict[str, Any]) -> bool:
        """验证输出格式"""
        required_fields = ["summary", "chart_type"]
        valid_chart_types = ["bar", "pie", "line", "scatter"]
        
        # 检查必需字段
        for field in required_fields:
            if field not in output:
                print(f"❌ 缺少必需字段: {field}")
                return False
        
        # 检查图表类型
        if output["chart_type"] not in valid_chart_types:
            print(f"❌ 无效的图表类型: {output['chart_type']}")
            return False
        
        return True
    
    def _fallback_format(self, formatter_input: FormatterInput) -> FormatterOutput:
        """备用格式转换逻辑"""
        try:
            print("🔄 使用备用格式转换逻辑...")
            
            query_results = formatter_input.query_results
            if not query_results:
                return FormatterOutput(
                    summary=formatter_input.summary,
                    chart_type="bar",
                    data=None
                )
            
            # 分析数据特征
            chart_type = self._analyze_chart_type(formatter_input.user_request, query_results)
            
            # 生成ECharts数据
            echarts_data = self._generate_echarts_data(query_results, chart_type)
            
            return FormatterOutput(
                summary=formatter_input.summary,
                chart_type=chart_type,
                data=echarts_data
            )
            
        except Exception as e:
            print(f"❌ 备用格式转换失败: {e}")
            return FormatterOutput(
                summary=formatter_input.summary,
                chart_type="bar",
                data=None
            )
    
    def _analyze_chart_type(self, user_request: str, query_results: List[Dict[str, Any]]) -> str:
        """分析并选择图表类型"""
        user_request_lower = user_request.lower()
        
        # 用户明确要求
        if "饼图" in user_request_lower or "占比" in user_request_lower:
            return "pie"
        elif "折线图" in user_request_lower or "趋势" in user_request_lower:
            return "line"
        elif "散点图" in user_request_lower or "关系" in user_request_lower:
            return "scatter"
        
        # 数据特征分析
        if not query_results:
            return "bar"
        
        first_row = query_results[0]
        data_count = len(query_results)
        
        # 检查字段类型
        numeric_fields = [k for k, v in first_row.items() if isinstance(v, (int, float))]
        text_fields = [k for k, v in first_row.items() if isinstance(v, str)]
        
        # 散点图：两个数值字段
        if len(numeric_fields) >= 2:
            return "scatter"
        
        # 饼图：数据量少且是占比类查询
        if data_count <= 8 and any(keyword in user_request_lower for keyword in ["占比", "分布", "比例"]):
            return "pie"
        
        # 折线图：时间相关
        if any(keyword in user_request_lower for keyword in ["趋势", "变化", "时间", "月份"]):
            return "line"
        
        # 默认柱状图
        return "bar"
    
    def _generate_echarts_data(self, query_results: List[Dict[str, Any]], chart_type: str) -> Optional[Dict[str, Any]]:
        """生成ECharts配置数据"""
        try:
            if not query_results:
                return None
            
            first_row = query_results[0]
            fields = list(first_row.keys())
            
            if chart_type == "bar":
                return self._generate_bar_data(query_results, fields)
            elif chart_type == "pie":
                return self._generate_pie_data(query_results, fields)
            elif chart_type == "line":
                return self._generate_line_data(query_results, fields)
            elif chart_type == "scatter":
                return self._generate_scatter_data(query_results, fields)
            else:
                return self._generate_bar_data(query_results, fields)
                
        except Exception as e:
            print(f"❌ ECharts数据生成失败: {e}")
            return None
    
    def _generate_bar_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成柱状图数据"""
        categories = [str(row.get(fields[0], "")) for row in query_results]
        values = [row.get(fields[1], 0) if len(fields) > 1 else 0 for row in query_results]
        
        return {
            "xAxis": {
                "type": "category",
                "data": categories
            },
            "yAxis": {
                "type": "value",
                "name": fields[1] if len(fields) > 1 else "数值"
            },
            "series": [{
                "name": fields[1] if len(fields) > 1 else "数值",
                "type": "bar",
                "data": values,
                "itemStyle": {"color": "#5470c6"}
            }]
        }
    
    def _generate_pie_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成饼图数据"""
        data = []
        for row in query_results:
            name = str(row.get(fields[0], ""))
            value = row.get(fields[1], 0) if len(fields) > 1 else 0
            data.append({"name": name, "value": value})
        
        return {
            "series": [{
                "name": fields[1] if len(fields) > 1 else "数值",
                "type": "pie",
                "radius": "50%",
                "data": data
            }]
        }
    
    def _generate_line_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成折线图数据"""
        categories = [str(row.get(fields[0], "")) for row in query_results]
        values = [row.get(fields[1], 0) if len(fields) > 1 else 0 for row in query_results]
        
        return {
            "xAxis": {
                "type": "category",
                "data": categories
            },
            "yAxis": {
                "type": "value",
                "name": fields[1] if len(fields) > 1 else "数值"
            },
            "series": [{
                "name": fields[1] if len(fields) > 1 else "数值",
                "type": "line",
                "data": values,
                "smooth": True
            }]
        }
    
    def _generate_scatter_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成散点图数据"""
        data = []
        for row in query_results:
            x = row.get(fields[0], 0)
            y = row.get(fields[1], 0) if len(fields) > 1 else 0
            data.append([x, y])
        
        return {
            "xAxis": {
                "type": "value",
                "name": fields[0]
            },
            "yAxis": {
                "type": "value",
                "name": fields[1] if len(fields) > 1 else "数值"
            },
            "series": [{
                "name": "数据点",
                "type": "scatter",
                "data": data,
                "symbolSize": 10
            }]
        }

# 全局实例
_formatter_agent = None

def get_formatter_agent(config=DS_CONFIG) -> DataFormatterAgent:
    """获取格式转换智能体实例"""
    global _formatter_agent
    if _formatter_agent is None:
        _formatter_agent = DataFormatterAgent(config)
    return _formatter_agent

def format_data_to_standard(user_request: str, query_results: List[Dict[str, Any]], summary: str, config=DS_CONFIG) -> Dict[str, Any]:
    """
    将数据分析结果转换为标准格式
    
    Args:
        user_request: 用户原始请求
        query_results: 数据库查询结果
        summary: 现有智能体的分析总结
        config: LLM配置
        
    Returns:
        Dict: 包含summary、chart_type和data的标准格式
    """
    formatter = get_formatter_agent(config)
    
    formatter_input = FormatterInput(
        user_request=user_request,
        query_results=query_results,
        summary=summary
    )
    
    result = formatter.format_data(formatter_input)
    
    return {
        "summary": result.summary,
        "chart_type": result.chart_type,
        "data": result.data
    }

if __name__ == "__main__":
    # 测试代码
    test_data = [
        {"ATTR_CITY": "朝阳区", "count": 185},
        {"ATTR_CITY": "丰台区", "count": 103},
        {"ATTR_CITY": "东城区", "count": 97}
    ]
    
    result = format_data_to_standard(
        user_request="请统计北京不同区中高端用户数",
        query_results=test_data,
        summary="北京市各区中高端用户数量统计完成"
    )
    
    print("测试结果:")
    print(json.dumps(result, ensure_ascii=False, indent=2))
